import pandas as pd
import torch
from fastapi import FastAPI, HTTPException

from models.ncf import NCF
from similarity import sql_similarity
import uvicorn
from pydantic import BaseModel
# 创建FastAPI应用
app = FastAPI(
    title="PY模块",
    version="1.0.0"
)

# 定义请求模型
class SQLCompareRequest(BaseModel):
    sql1: str
    sql2: str

# 定义响应模型
class SQLSimilarityResponse(BaseModel):
    similarity: float
    sql1: str
    sql2: str

# 加载模型和映射
df = pd.read_csv("data/ratings.csv")
# 如果训练时使用的是这种方式
num_users = df['user_id'].max() + 1
num_items = df['course_id'].max() + 1

model = NCF(num_users, num_items, mlp_layers=[32, 16, 8])
model.load_state_dict(torch.load("ncf_model.pth"))
model.eval()

# 请求体模型
class RecommendationRequest(BaseModel):
    user_id: int
    top_k: int = 5

@app.post("/recommend")
def recommend(request: RecommendationRequest):
    user_id = request.user_id
    top_k = request.top_k

    items = torch.arange(1, num_items)
    users = torch.full_like(items, fill_value=user_id)

    with torch.no_grad():
        scores = model(users, items).numpy()

    top_items = items.numpy()[scores.argsort()[::-1][:top_k]]
    return {"recommended_courses": list(map(int, top_items))}

@app.post("/compare-sql", response_model=SQLSimilarityResponse)
async def compare_sql(request: SQLCompareRequest):
    """计算两个SQL语句的重复度
    
    Args:
        request: 包含两个SQL语句的请求
        
    Returns:
        包含重复度(相似度)的响应
    """
    try:
        # 计算相似度
        similarity = sql_similarity(request.sql1, request.sql2)
        
        # 返回结果
        return SQLSimilarityResponse(
            similarity=similarity,
            sql1=request.sql1,
            sql2=request.sql2
        )
    except Exception as e:
        raise HTTPException(status_code=500, detail=f"计算相似度时出错: {str(e)}")

# 添加根路径响应
@app.get("/")
async def root():
    return {
        "message": "SQL相似度API服务正在运行，请使用 /compare-sql 接口"}

# 启动服务器的代码
if __name__ == "__main__":
    uvicorn.run(app, host="0.0.0.0", port=8000)